package com.suemi.missile.common.mql.parse;


import com.suemi.missile.common.mql.domain.*;
import com.suemi.missile.common.mql.exception.MqlParseException;
import org.antlr.v4.runtime.ANTLRInputStream;
import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.tree.ParseTree;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@SuppressWarnings("Unchecked")
public class MqlParseVisitor extends MqlBaseVisitor<Object> {

  public static Query generate(String mql) throws MqlParseException {
    CharStream input = new ANTLRInputStream(mql);
    MqlLexer lexer = new MqlLexer(input);
    lexer.removeErrorListeners();
    lexer.addErrorListener(ParseErrorListener.instance);
    CommonTokenStream tokens = new CommonTokenStream(lexer);
    MqlParser parser = new MqlParser(tokens);
    parser.removeErrorListeners();
    parser.addErrorListener(ParseErrorListener.instance);

    MqlParseVisitor visitor = new MqlParseVisitor();

    return visitor.parse(parser.query());
  }

  public Query parse(MqlParser.QueryContext ctx) throws MqlParseException {
    try {
      return (Query) visitQuery(ctx);
    } catch (MqlParseException e) {
      throw e;
    } catch (Exception e) {
      throw new MqlParseException(e);
    }
  }

  @Override
  public Object visitQuery(MqlParser.QueryContext ctx) {
    return visit(ctx.select());
  }

  @Override
  public Object visitSelect(MqlParser.SelectContext ctx) {
    List<Field> fields = (List<Field>) visitFieldList(ctx.fieldList());
    Pair<Integer, Integer> limit = (Pair<Integer, Integer>) visitLimit(ctx.limit());
    List<Order> orders = (List<Order>) visitOrderbyList(ctx.orderbyList());
    From from = (From) visitFrom(ctx.from());
    Where where = (Where) visitWhere(ctx.where());
    Select select = new Select();
    fields.forEach(select::addField);
    select.addFrom(from);
    if (orders != null) orders.forEach(select::addOrderBy);
    select.setWhere(where);
    if (limit != null) {
      select.setOffset(limit.getLeft());
      select.setRowCount(limit.getRight());
    }
    return select;
  }

  @Override
  public Object visitFrom(MqlParser.FromContext ctx) {
    if(ctx == null) return null;
    return ctx.ID().size() == 1 ? new From(ctx.ID(0).getText()) : new From(String.join("/", ctx.ID().stream().map(ParseTree::getText).collect(Collectors.toList())));
  }

  @Override
  public Object visitOrderby(MqlParser.OrderbyContext ctx) {
    if(ctx == null) return null;
    return ctx.getChildCount() > 1 ? new Order(ctx.ID().getText(), ctx.getChild(1).getText()) : new Order(ctx.ID().getText(), "ASC");
  }

  @Override
  public Object visitOrderbyList(MqlParser.OrderbyListContext ctx) {
    if(ctx == null) return null;
    List<Order> orders = new ArrayList<>();
    ctx.orderby().forEach(c -> {
      orders.add((Order) visit(c));
    });
    return orders;
  }

  @Override
  public Object visitLimit(MqlParser.LimitContext ctx) {
    if(ctx == null) return null;
    return ctx.INT().size() == 1 ? new ImmutablePair<Integer, Integer>(0, Integer.parseInt(ctx.INT(0).getText())) :
            new ImmutablePair<Integer, Integer>(Integer.parseInt(ctx.INT(0).getText()), Integer.parseInt(ctx.INT(1).getText()));
  }

  @Override
  public Object visitFieldList(MqlParser.FieldListContext ctx) {
    if(ctx == null) return null;
    List<Field> fields = new ArrayList<>();
    if (ctx.getText().equals("*"))
      fields.add(new Field("*", null));
    else
      ctx.ID().forEach(id -> fields.add(new Field(id.getText(), null)));
    return fields;
  }

  @Override
  public Object visitWhere(MqlParser.WhereContext ctx) {
    return visit(ctx.cond());
  }


  @Override
  public Object visitParCond(MqlParser.ParCondContext ctx) {
    return visit(ctx.cond());
  }

  @Override
  public Object visitOrCond(MqlParser.OrCondContext ctx) {
    Where where = new Where(Where.CONN.OR);
    ctx.cond().forEach(c-> {
      Where cond = (Where) visit(c);
      if(cond == null) return;
//      if(cond.getBoost()!=null) {
//        where.setBoost(where.getBoost()==null?cond.getBoost():Math.max(where.getBoost(),cond.getBoost()));
//      }
      where.addWhere(cond);
    });
    return where;
  }

  @Override
  public Object visitAndCond(MqlParser.AndCondContext ctx) {
    Where where = new Where(Where.CONN.AND);
    ctx.cond().forEach(c-> {
      Where cond = (Where) visit(c);
      if(cond == null) return;
//      if(cond.getBoost()!=null) {
//        where.setBoost(where.getBoost()==null?cond.getBoost():Math.max(where.getBoost(),cond.getBoost()));
//      }
      where.addWhere(cond);
    });
    return where;
  }

  @Override
  public Object visitNotCond(MqlParser.NotCondContext ctx) {
    Where where = new Where(Where.CONN.NOT);
    Where cond = (Where) visit(ctx.cond());
    where.setBoost(cond.getBoost());
    where.addWhere(cond);
    return where;
  }

  @Override
  public Object visitContainsCond(MqlParser.ContainsCondContext ctx) {
    FieldTreeVisitor visitor = new FieldTreeVisitor(ctx.wordExpr(),ctx);
    return visitor.visit(ctx.fieldExpr());
  }

}
